import argparse
import os
import torch
from sympy.core.random import shuffle

from lib.weld_data import get_dataloader
from lib.model import Ganomaly
from lib.utils import *


def train_ganomaly(opt):
    """
    训练GANomaly模型，适配文件夹结构：
    根目录/
        train/
            normal_regions/  # 正常样本（必选）
            defect_regions/  # 异常样本（仅半监督模式需要）
            masks/           # 可选，掩码文件
        test/
            normal_regions/  # 测试集正常样本（必选）
            defect_regions/  # 测试集异常样本（必选）
            masks/           # 可选，掩码文件
    """
    # 1. 数据路径校验（根据模式动态调整）
    required_dirs = [
        os.path.join(opt.dataroot, 'train', 'normal_regions'),  # 训练集正常样本必选
        os.path.join(opt.dataroot, 'test', 'normal_regions'),  # 测试集正常样本必选
        os.path.join(opt.dataroot, 'test', 'defect_regions')  # 测试集异常样本必选
    ]
    # 半监督模式下，额外校验训练集异常样本目录
    if opt.use_semi:
        required_dirs.append(os.path.join(opt.dataroot, 'train', 'defect_regions'))

    # 检查所有必要目录是否存在
    for dir_path in required_dirs:
        if not os.path.exists(dir_path):
            raise FileNotFoundError(
                f"数据集路径不存在: {dir_path}\n请检查文件夹结构是否正确（当前模式：{'半监督' if opt.use_semi else '无监督'}）")

    # 2. 加载数据
    print(f"==> 从 {opt.dataroot} 加载数据...（模式：{'半监督' if opt.use_semi else '无监督'}）")
    dataloader = {
        'train': get_dataloader(
            root=opt.dataroot,
            batch_size=opt.batchsize,
            workers=opt.workers,
            isize=opt.isize,
            phase='train',
            include_train_defect=opt.use_semi,  # 半监督模式加载训练集异常样本
            # droplast=opt.droplast
        ),
        'test': get_dataloader(
            root=opt.dataroot,
            batch_size=opt.batchsize,
            workers=opt.workers,
            isize=opt.isize,
            phase='test',
            # droplast=opt.droplast
        )
    }

    # 3. 打印数据集详细信息（增强可解释性）
    # 训练集正常/异常样本数量（半监督模式下）
    train_normal_cnt = len([f for f in os.listdir(os.path.join(opt.dataroot, 'train', 'normal_regions')) if
                            f.endswith(('.png', '.jpg', '.jpeg'))])
    train_anomaly_cnt = 0
    if opt.use_semi:
        train_anomaly_cnt = len([f for f in os.listdir(os.path.join(opt.dataroot, 'train', 'defect_regions')) if
                                 f.endswith(('.png', '.jpg', '.jpeg'))])

    # 测试集正常/异常样本数量
    test_normal_cnt = len([f for f in os.listdir(os.path.join(opt.dataroot, 'test', 'normal_regions')) if
                           f.endswith(('.png', '.jpg', '.jpeg'))])
    test_anomaly_cnt = len([f for f in os.listdir(os.path.join(opt.dataroot, 'test', 'defect_regions')) if
                            f.endswith(('.png', '.jpg', '.jpeg'))])

    print(f"==> 数据加载完成:")
    print(
        f"    训练集 | 总样本数: {len(dataloader['train'].dataset)} | 正常样本: {train_normal_cnt} | 异常样本: {train_anomaly_cnt}")
    print(
        f"    测试集 | 总样本数: {len(dataloader['test'].dataset)} | 正常样本: {test_normal_cnt} | 异常样本: {test_anomaly_cnt}")

    # 4. 设备校验与配置（保持原有逻辑，优化打印信息）
    str_ids = opt.gpu_ids.split(',')
    opt.gpu_ids = []
    for str_id in str_ids:
        id = int(str_id)
        if id >= 0:
            opt.gpu_ids.append(id)
    # 设备选择逻辑
    if opt.device == 'gpu' and len(opt.gpu_ids) > 0 and torch.cuda.is_available():
        opt.device = torch.device(f"cuda:{opt.gpu_ids[0]}")
        print(f"==> 使用GPU设备: {opt.device} (ID: {opt.gpu_ids})")
    else:
        opt.device = torch.device('cpu')
        print(f"==> 使用CPU设备 (GPU不可用或未指定)")

    # 5. 创建输出目录（补充可解释性相关目录）
    # 核心目录
    os.makedirs(os.path.join(opt.outf, opt.name, 'train', 'weights'), exist_ok=True)  # 模型权重
    os.makedirs(os.path.join(opt.outf, opt.name, 'test', 'images'), exist_ok=True)  # 测试图像
    # 可解释性相关目录
    os.makedirs(os.path.join(opt.outf, opt.name, 'train', 'reconstructions'), exist_ok=True)  # 训练重建对比
    os.makedirs(os.path.join(opt.outf, opt.name, 'test', 'interpretability'), exist_ok=True)  # 测试可解释性结果
    if opt.save_test_images:
        os.makedirs(os.path.join(opt.outf, opt.name, 'test', 'demo'), exist_ok=True)  # 测试展示图像

    # 6. 初始化模型并训练
    print("==> 初始化GANomaly模型...")
    # 半监督模式下提示分类损失启用
    if opt.use_semi:
        print(f"==> 半监督模式启用 | 分类损失权重: {opt.w_cls}")
    if opt.use_contrast:
        print(f"==> 对比学习启用 | 对比损失权重: {opt.w_contrast}")
    if opt.use_push:
        print(f"==>异常推开损失启用 | 推开损失权重: {opt.w_push}")
    if opt.use_latent_reg:
        print(f"==>L2潜在向量正则化启用 | 潜在正则化损失权重: {opt.w_latent_reg}")
    config_dir = os.path.join(opt.outf, opt.name, 'configs')
    save_experiment_config(opt, config_dir)

    model = Ganomaly(opt, dataloader)

    print("==> 开始训练...")
    model.train()

    print(f"==> 训练完成！模型权重保存在：{os.path.join(opt.outf, opt.name, 'train', 'weights')}")
    print(f"==> 可解释性结果保存在：{os.path.join(opt.outf, opt.name, 'test', 'interpretability')}")


if __name__ == '__main__':
    # 使用带默认值说明的格式化器
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # 基础参数
    parser.add_argument('--dataset', default='folder', help='数据集类型 (folder | cifar10 | mnist)')
    parser.add_argument('--dataroot', default='./data/weld_dataset', help='数据集根目录')
    parser.add_argument('--batchsize', type=int, default=32, help='输入批次大小')
    parser.add_argument('--workers', type=int, default=0, help='数据加载线程数 (Windows建议设为0)')
    parser.add_argument('--droplast', action='store_true', default=True, help='是否丢弃最后不完整的批次')
    parser.add_argument('--isize', type=int, default=64, help='输入图像尺寸 (宽=高)')
    parser.add_argument('--nc', type=int, default=3, help='输入图像通道数 (RGB=3, 灰度=1)')
    parser.add_argument('--nz', type=int, default=100, help='潜在向量z的维度')
    parser.add_argument('--ngf', type=int, default=64, help='生成器特征图数量基数')
    parser.add_argument('--ndf', type=int, default=64, help='判别器特征图数量基数')
    parser.add_argument('--extralayers', type=int, default=0, help='生成器和判别器的额外层数量')
    parser.add_argument('--device', type=str, default='gpu', help='计算设备 (gpu | cpu)')
    parser.add_argument('--gpu_ids', type=str, default='0', help='GPU设备ID (例如 0 或 0,1 使用-1表示CPU)')
    parser.add_argument('--ngpu', type=int, default=1, help='使用的GPU数量')
    parser.add_argument('--name', type=str, default='ganomaly_weld_semi_supervised_contrast_push_hlfd_1_50_10_v1', help='实验名称')
    parser.add_argument('--model', type=str, default='ganomaly', help='使用的模型 (目前仅支持ganomaly)')
    parser.add_argument('--display_server', type=str, default="http://localhost", help='Visdom服务器地址')
    parser.add_argument('--display_port', type=int, default=8097, help='Visdom服务器端口')
    parser.add_argument('--display_id', type=int, default=0, help='Visdom显示窗口ID')
    parser.add_argument('--display', action='store_true', help='是否使用Visdom可视化')
    parser.add_argument('--outf', default='./output', help='输出图像和模型 checkpoint 的文件夹')
    parser.add_argument('--manualseed', default=-1, type=int, help='手动随机种子 (-1表示不固定)')
    parser.add_argument('--abnormal_class', default='defect_regions', help='异常类名称 (针对焊缝数据集)')
    parser.add_argument('--proportion', type=float, default=0.3, help='测试集中异常样本的比例')
    parser.add_argument('--metric', type=str, default='roc', help='评估指标 (roc | auc)')

    # 训练参数
    parser.add_argument('--print_freq', type=int, default=5, help='控制台打印训练结果的频率')
    parser.add_argument('--save_image_freq', type=int, default=20, help='保存真实和生成图像的频率')
    parser.add_argument('--resume', default='', help='用于继续训练的checkpoint路径')
    parser.add_argument('--phase', type=str, default='train', help='运行阶段 (train | val | test)')
    parser.add_argument('--iter', type=int, default=0, help='开始训练的迭代次数')
    parser.add_argument('--niter', type=int, default=200, help='训练的总轮数')
    parser.add_argument('--beta1', type=float, default=0.5, help='Adam优化器的momentum参数')
    parser.add_argument('--lr', type=float, default=0.0002, help='初始学习率')



    parser.add_argument('--w_adv', type=float, default=1.0, help='对抗损失权重')
    parser.add_argument('--w_con', type=float, default=50.0, help='重建损失权重')
    parser.add_argument('--w_enc', type=float, default=10.0, help='编码器一致性损失权重')
    parser.add_argument('--w_cls', type=float, default=5.0, help='异常分类损失的权重')
    parser.add_argument('--w_contrast', type=float, default=1.0, help='对比学习损失权重')
    parser.add_argument('--w_push', type=float, default=10.0, help='异常推开损失权重')
    parser.add_argument('--w_latent_reg', type=float, default=1.0, help='潜在正则化损失权重')
    parser.add_argument('--contrast_temp', type=float, default=0.05, help='对比学习温度参数')


    parser.add_argument('--vis_freq', type=int, default=20, help='可视化结果保存频率（每多少轮保存一次）')
    parser.add_argument('--save_latent', action='store_true', help='是否保存潜在向量用于TSNE（会增加内存占用）')
    parser.add_argument('--use_contrast', action='store_true', help='启用对比学习损失')
    parser.add_argument('--use_semi', action='store_true', help='是否启用半监督模式（训练集含异常样本）')
    parser.add_argument('--save_test_images', action='store_true', help='是否保存测试图像用于展示')
    parser.add_argument('--load_weights', action='store_true', help='是否加载预训练权重')
    parser.add_argument('--use_push', action='store_true', help='启用异常推开损失')
    parser.add_argument('--use_latent_reg', action='store_true', help='启用潜在向量正则化')
    parser.add_argument('--use_hlfd', action='store_true', help='启用潜在向量正则化')
    parser.add_argument('--use_mask', action='store_true', help='启用潜在向量正则化')

    opt = parser.parse_args()
    if opt.phase == 'train':
        opt.isTrain = True
    else:
        opt.isTrain = False
    train_ganomaly(opt)